#!/usr/bin/env python
# encoding: utf-8
'''
@author: wangjianrong
@software: pycharm
@file: cuda_info.py
@time: 2020/9/23 15:46
@desc:
'''

import torch


#获得gpu数量
device_cnt = torch.cuda.device_count()
device_ids = list(range(device_cnt))
#获得gpu属性
for gpu_id in device_ids:
    device_prop = torch.cuda.get_device_properties(gpu_id)
    device_cap = torch.cuda.get_device_capability(gpu_id)
    print(device_prop)
    print(device_cap)